#include <src/pmc.h>
#include <src/pmc_walker.h>
#include <src/random.h>
#include <src/timing.h>
#include <src/move.h>

//--------------------------------------------------------------------------------
// pcps::pmc::split_and_combine -- splits and combines walkers to keep their
//                                 weights close to one
//--------------------------------------------------------------------------------

template <class SCALAR, class PART, class REF, class GENERATOR>
void pcps::pmc::split_and_combine(const pcps::Input & userinp,
                                  const pcps::Hamiltonian<SCALAR> & ham,
                                  const pcps::Wavefunction<SCALAR> & wfn,
                                  const double n_active,
                                  const bool fixed_phase,
                                  pcps::Arena & arena,
                                  GENERATOR & gen,
                                  std::vector<boost::shared_ptr<pcps::pmc::Walker<SCALAR, PART, REF> > > & walkers,
                                  std::list<boost::shared_ptr<pcps::pmc::Walker<SCALAR, PART, REF> > > & active_walkers,
                                  std::list<boost::shared_ptr<pcps::pmc::Walker<SCALAR, PART, REF> > > & inactive_walkers) {

  // get MPI info
  const MPI::Comm & comm = MPI::COMM_WORLD;
  const int nproc = comm.Get_size();
  const int myrank = comm.Get_rank();

  // create a short name for a shared pointer to a walker
  typedef boost::shared_ptr<pcps::pmc::Walker<SCALAR, PART, REF> > WalkerPtr;

  // create a list to hold walkers with small weights
  std::list<WalkerPtr> small_walkers;

  // split large walkers and move small walkers to a list so they can be combined
  for (typename std::list<WalkerPtr>::iterator aw = active_walkers.begin(); aw != active_walkers.end(); ) {

    // check that the walker's weight is sane
    assert( !fixed_phase || (**aw).weight() >= 0.0 );

    // split the walker if its weight is too large
    if (std::abs((**aw).weight()) > 2.0) {

      // determine how many copies to create
      const int n_copies = int(std::abs((**aw).weight())) - 1;

      // adjust the weight by the new number of walkers
      (**aw).weight() /= n_copies + 1;

      // create the copies
      for (int i = 0; i < n_copies; i++) {

        // if we have run out of inactive walkers, create a new one
        if ( inactive_walkers.empty() ) {
          walkers.push_back( WalkerPtr( new pcps::pmc::Walker<SCALAR, PART, REF>(userinp, ham, wfn, arena) ) );
          inactive_walkers.push_back(*walkers.rbegin());
        }

        // copy the walker
        (**inactive_walkers.begin()).copy(userinp, ham, **aw);

        // add the copy to the active walker list
        active_walkers.insert(aw, *inactive_walkers.begin());

        // remove the copy from the inactive walker list
        inactive_walkers.pop_front();

      }

      // move to the next walker
      aw++;

    // move walkers with small weights to a temporary list to be combined later
    } else if ( std::abs((**aw).weight()) < 0.5 ) {

      // add the walker to the list of small walkers
      small_walkers.push_back(*aw);

      // remove the walker from the active walker list and move to the next walker
      aw = active_walkers.erase(aw);

    // leave walkers with reasonable weights alone
    } else {

      // move to the next walker
      aw++;

    }

  }

  // combine small walkers
  if (fixed_phase) {

    while (small_walkers.size() > 0) {

      // demarcate a group of walkers with net weight >= 1, or all the remaining small walkers if their total weight is < 1
      double group_weight = 0.0;
      typename std::list<WalkerPtr>::iterator group_end;
      for ( group_end = small_walkers.begin(); group_weight < 1.0 && group_end != small_walkers.end(); group_end++) {
        assert( (**group_end).weight() >= 0.0 );
        group_weight += (**group_end).weight();
      }

      // pick a randon number uniformly between zero and one to choose the walker that will survive
      double r = pcps::uni_01(gen);

      // find the chosen walker
      while (small_walkers.begin() != group_end) {

        r -= (**small_walkers.begin()).weight() / group_weight;

        // stop at the chosen walker
        if (r < 1.0e-10) break;

        // if this walker was not chosen, move it to the inactive list
        inactive_walkers.push_front(*small_walkers.begin());
        small_walkers.pop_front();

      }

      // raise an error if we did not choose any of the walkers in the group
      if ( small_walkers.begin() == group_end )
        throw pcps::Exception( (boost::format("problem combining small walkers in pmc population control")).str() );

      // give the chosen walker the entire group's weight
      (**small_walkers.begin()).weight() = group_weight;

      // move the chosen walker to the active walker list
      active_walkers.push_back(*small_walkers.begin());
      small_walkers.pop_front();

      // move any remaining walkers in the group to the inactive list
      while (small_walkers.begin() != group_end) {
        inactive_walkers.push_front(*small_walkers.begin());
        small_walkers.pop_front();
      }

    }

  // for walkers with positive or negative weights
  } else {

    while (small_walkers.size() > 0) {
      if (pcps::uni_01(gen) < std::abs((**small_walkers.begin()).weight()) ) {
        (**small_walkers.begin()).weight() /= std::abs((**small_walkers.begin()).weight());
        active_walkers.push_back(*small_walkers.begin());
      } else {
        inactive_walkers.push_front(*small_walkers.begin());
      }
      small_walkers.pop_front();
    }

  }

}

//--------------------------------------------------------------------------------
// pcps::pmc::driver_detail -- performs a projector monte carlo calculation
//--------------------------------------------------------------------------------

template <class SCALAR, class PART, class REF>
void pcps::pmc::driver_detail(const pcps::Input & userinp, const pcps::Hamiltonian<SCALAR> & ham, const pcps::Wavefunction<SCALAR> & wfn) {

  // get MPI info
  const MPI::Comm & comm = MPI::COMM_WORLD;
  const int nproc = comm.Get_size();
  const int myrank = comm.Get_rank();

  // ensure that the user has not requested multiple threads
  if (userinp.nthreads() != 1)
    throw pcps::Exception( (boost::format("projected monte carlo is not multi-threaded, so nthreads must be 1")).str() );

  // create an arena object for memory allocation
  pcps::Arena arena(userinp.arena_block_size(), userinp.arena_buffer_size());

  // create a random number generator
  pcps::LaggedFibonacci<1279,418> gen(arena);

  // create a short name for a shared pointer to a walker
  typedef boost::shared_ptr<pcps::pmc::Walker<SCALAR, PART, REF> > WalkerPtr;

  // create a vector to hold this process's walkers
  std::vector<WalkerPtr> walkers( 2 * userinp.walkers_per_process() );

  // create the walkers
  for (int i = 0; i < walkers.size(); i++)
    walkers.at(i) = WalkerPtr( new pcps::pmc::Walker<SCALAR, PART, REF>(userinp, ham, wfn, arena) );

  // create a list of active walkers
  std::list<WalkerPtr> active_walkers;
  active_walkers.push_back(walkers.at(0));
  //for (int i = 0; i < userinp.walkers_per_process(); i++)
  //  active_walkers.push_back(walkers.at(i));

  // create a list of inactive walkers
  std::list<WalkerPtr> inactive_walkers;
  for (int i = 1; i < walkers.size(); i++)
    inactive_walkers.push_back(walkers.at(i));
  //for (int i = userinp.walkers_per_process(); i < walkers.size(); i++)
  //  inactive_walkers.push_back(walkers.at(i));

  // assign starting configurations to the active walkers
  pcps::MoveFunc<SCALAR, PART, REF> mv(userinp, ham, wfn, arena);
  for (typename std::list<WalkerPtr>::iterator aw = active_walkers.begin(); aw != active_walkers.end(); aw++) {
    (**aw).init(userinp, ham, mv, gen);
  }

  // create variables for tracking the moving average
  double moving_energy = 0.0;
  double moving_weight = 0.0;
  std::list<double> energy_list, weight_list;
  if (userinp.print_freq() != userinp.pop_control_freq())
    throw pcps::Exception( (boost::format("pmc moving averages currently require print_freq and pop_control_freq to be equal")).str() );

  // iterate
  assert( userinp.sample_length() > 0 );
  for (long int iter = 0; iter < userinp.sample_length(); iter++) {

    // determine whether to fix the nodes for this iteration
    const bool fix_nodes = ( userinp.begin_true_projector() <= 0 || iter < userinp.begin_true_projector() );

    // compute the energy, total weight, and number of active walkers
    double energy = 0.0;
    double total_weight = 0.0;
    double n_active = active_walkers.size();
    for (typename std::list<WalkerPtr>::iterator aw = active_walkers.begin(); aw != active_walkers.end(); aw++) {
      energy += (**aw).weight() * pcps::real((**aw).local_energy());
      total_weight += (**aw).weight();
    }

    // update moving totals
    moving_energy += energy;
    moving_weight += total_weight;
    energy_list.push_back(energy);
    weight_list.push_back(total_weight);
    if (energy_list.size() > userinp.print_freq()) {
      moving_energy -= *energy_list.begin();
      moving_weight -= *weight_list.begin();
      energy_list.pop_front();
      weight_list.pop_front();
    }

    // reduce energy, weight, and number of active walkers to all processes
    if ( (iter+1) % userinp.print_freq() == 0 || (iter+1) % userinp.pop_control_freq() == 0 ) {

      // prepare vectors for sending/receiving data
      std::vector<double> send_buf(3, 0.0);
      std::vector<double> recv_buf(3, 0.0);
      send_buf.at(0) = moving_energy / energy_list.size();
      send_buf.at(1) = moving_weight / weight_list.size();
      send_buf.at(2) = n_active;

      // reduce the data
      pcps::allreduce(comm, &send_buf.at(0), &recv_buf.at(0), 3, MPI::SUM);

      // compute total energy, total weight, and total number of active walkers
      energy = recv_buf.at(0) / recv_buf.at(1);
      total_weight = recv_buf.at(1);
      n_active = recv_buf.at(2);

    }

    // print energy, total weight, and number of walkers on root process
    if (myrank == 0 && (iter+1) % userinp.print_freq() == 0)
      std::cout << boost::format("iteration %22i                   elapsed time = %17.6f sec")
                   % (iter+1)
                   % pcps::get_overall_timer().elapsed_seconds()
                << std::endl
                << boost::format("energy                 %20.12f")
                   % energy
                << std::endl
                << boost::format("total weight           %20.12f")
                   % total_weight
                << std::endl
                << boost::format("active walkers %17i")
                   % n_active
                << std::endl
                << std::endl;

    // every so many iterations, perform population control by adjusting each walker's weight
    if ( (iter+1) % userinp.pop_control_freq() == 0 )
      for (typename std::list<WalkerPtr>::iterator aw = active_walkers.begin(); aw != active_walkers.end(); aw++)
        (**aw).weight() *= nproc * double(userinp.walkers_per_process()) / n_active;

    // split large walkers and combine small walkers to keep the walkers' weights close to one
    pcps::pmc::split_and_combine(userinp, ham, wfn, n_active, fix_nodes, arena, gen, walkers, active_walkers, inactive_walkers);

    // move the walkers and adjust their weights according to the application of the transition matrix
    int walker_number = 0;
    for (typename std::list<WalkerPtr>::iterator aw = active_walkers.begin(); aw != active_walkers.end(); aw++) {
      //std::cout << boost::format("attempting to move walker %4i") % (walker_number++) << std::endl;
      (**aw).move(userinp, ham, gen, fix_nodes);
    }

  }

}

//--------------------------------------------------------------------------------
// pcps::pmc::driver -- performs a projector monte carlo calculation
//--------------------------------------------------------------------------------

template <class SCALAR>
void pcps::pmc::driver(const pcps::Input & userinp, const pcps::Hamiltonian<SCALAR> & ham, const pcps::Wavefunction<SCALAR> & wfn) {

  // spins with the uniform reference
  if (userinp.particle_type() == pcps::keyword::Part_Spin && userinp.ref_type() == pcps::keyword::Ref_Uniform)
    pcps::pmc::driver_detail<SCALAR, pcps::Part_Spin, pcps::Ref_Uniform>(userinp, ham, wfn);

  // spins with the AGP reference
  else if ( userinp.particle_type() == pcps::keyword::Part_Spin && userinp.ref_type() == pcps::keyword::Ref_AGP )
    pcps::pmc::driver_detail<SCALAR, pcps::Part_Spin, pcps::Ref_AGP>(userinp, ham, wfn);

  // spins with the pairing reference
  else if ( userinp.particle_type() == pcps::keyword::Part_Spin && userinp.ref_type() == pcps::keyword::Ref_Pairing )
    pcps::pmc::driver_detail<SCALAR, pcps::Part_Spin, pcps::Ref_Pairing>(userinp, ham, wfn);

  // fermions with the uniform reference
  else if (userinp.particle_type() == pcps::keyword::Part_Fermion && userinp.ref_type() == pcps::keyword::Ref_Uniform)
    pcps::pmc::driver_detail<SCALAR, pcps::Part_Fermion, pcps::Ref_Uniform>(userinp, ham, wfn);

  // fermions with the restricted Hartree Fock or AGP reference
  else if (   userinp.particle_type() == pcps::keyword::Part_Fermion
           && ( userinp.ref_type() == pcps::keyword::Ref_RHF || userinp.ref_type() == pcps::keyword::Ref_AGP ) )
    pcps::pmc::driver_detail<SCALAR, pcps::Part_Fermion, pcps::Ref_AGP>(userinp, ham, wfn);

  // fermions with the pairing reference
  else if ( userinp.particle_type() == pcps::keyword::Part_Fermion && userinp.ref_type() == pcps::keyword::Ref_Pairing )
    pcps::pmc::driver_detail<SCALAR, pcps::Part_Fermion, pcps::Ref_Pairing>(userinp, ham, wfn);

  // fermions with the pfaffian reference
  else if ( userinp.particle_type() == pcps::keyword::Part_Fermion && userinp.ref_type() == pcps::keyword::Ref_Pfaffian )
    pcps::pmc::driver_detail<SCALAR, pcps::Part_Fermion, pcps::Ref_Pfaffian>(userinp, ham, wfn);

  // unsupported combination of particle and reference type
  else
    throw pcps::Exception(
      (boost::format("Error in pcps::pmc::driver.  The combination of particle type \"%s\" and reference type \"%s\" is not supported.")
        % userinp.particle_type()
        % userinp.ref_type()
      ).str()
    );

}

// explicitly instantiate the template
template void pcps::pmc::driver(const pcps::Input &,
                                const pcps::Hamiltonian<double> &,
                                const pcps::Wavefunction<double> &);
template void pcps::pmc::driver(const pcps::Input &,
                                const pcps::Hamiltonian<std::complex<double> > &,
                                const pcps::Wavefunction<std::complex<double> > &);
